import torch
import torch.nn as nn
import copy
import random
import sys
from tqdm import tqdm
from torch.utils.data.dataset import Dataset
import os
import numpy as np
import csv

# 用来记录扰动过程的记录器
class Disturb_Recorder():
    def __init__(self, file_path: str) -> None:
        self.file_path = file_path
        self.f = open(file_path, 'w', newline='')
        self.writer = csv.writer(self.f)
        # 写入表头
        self.writer.writerow([
            "Rank",
            "Original Code",
            "Disturb Code",
            "True label"
        ])
        # 非常抱歉，这个"True label"的label是小写的，但是我发现的时候，已经用很多了，我不想改了

    def write(self, rank: int, original_code: str, disturbed_code: str, label: str) -> None:
        """写入一行记录"""
        self.writer.writerow([
            rank,
            original_code,
            disturbed_code,
            label
        ])

    def close(self):
        """关闭文件"""
        self.f.close()

    # 使用上下文管理器自动关闭文件
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


# 用来记录扰动结果的记录器
class Adversarial_Recorder():
    def __init__(self, file_path: str) -> None:
        self.file_path = file_path
        self.f = open(file_path, 'w', newline='')
        self.writer = csv.writer(self.f)
        # 写入表头
        self.writer.writerow([
            "Index",
            "Original Code",
            "Adversarial Code",
            "True Label",
            "Original Prediction",
            "Adv Prediction"
        ])

    def write(self, index: int, original_code: str, adv_code: str, true_label: str, orig_pred: str, adv_pred:str) -> None:
        """写入一行记录"""
        self.writer.writerow([
            index,
            original_code,
            adv_code,
            true_label,
            orig_pred,
            adv_pred
        ])

    def close(self):
        """关闭文件"""
        self.f.close()

    # 使用上下文管理器自动关闭文件
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()